import torch
import torch.nn as nn
import torch.nn.functional as F
from config import Config
# from pytorch_pretrained_bert import BertModel, BertTokenizer, BertConfig, BertForSequenceClassification
from pytorch_bert import BertModel, BertTokenizer, BertConfig, BertForSequenceClassification



class BertClassifier(nn.Module):
    def __init__(self, config):
        super(BertClassifier, self).__init__()
        self.bert = BertForSequenceClassification.from_pretrained(config.bert_path, num_labels=config.num_classes)

    def forward(self, data, noise=None):
        logits = self.bert(input_ids=data[0],
                           attention_mask=data[3],
                           labels=None,
                           noise=noise)
        return logits

